import os
import json
import glob
from collections import defaultdict

def fix_annotations_structure(project_dir):
    """
    修复标注文件结构，将混合在一起的TI和RGB标注按模态分离到对应文件夹
    
    Args:
        project_dir: 项目根目录路径
        
    Returns:
        dict: 处理结果统计信息
    """
    frames_dir = os.path.join(project_dir, "Frames")
    if not os.path.exists(frames_dir):
        return {"error": f"找不到Frames目录: {frames_dir}"}
        
    # 扫描所有可能包含标注的目录
    result = {
        "processed": 0,
        "fixed": 0,
        "errors": [],
        "details": []
    }
    
    for root, dirs, files in os.walk(frames_dir):
        # 检查当前目录是否包含T或V子目录
        has_t = "T" in dirs
        has_v = "V" in dirs
        
        # 只处理包含T或V子目录的目录
        if not (has_t or has_v):
            continue
            
        # 检查各处可能存在的标注文件
        source_annotations = None
        t_folder = os.path.join(root, "T") if has_t else None
        v_folder = os.path.join(root, "V") if has_v else None
        
        # 查找source_annotations (不管在哪里找到都可以)
        possible_paths = [
            os.path.join(root, "annotations.json"),  # 根目录
            os.path.join(t_folder, "annotations.json") if t_folder else None,  # T文件夹
            os.path.join(v_folder, "annotations.json") if v_folder else None,  # V文件夹
        ]
        
        for path in possible_paths:
            if path and os.path.exists(path):
                try:
                    with open(path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        if "annotations" in data and isinstance(data["annotations"], list):
                            source_annotations = data
                            source_path = path
                            break
                except Exception as e:
                    result["errors"].append(f"读取标注文件失败: {path}, 错误: {str(e)}")
                    continue
        
        if not source_annotations:
            continue  # 没有找到有效的标注文件
            
        result["processed"] += 1
        
        # 按模态分离标注
        annotations = source_annotations.get("annotations", [])
        rgb_annotations = [ann for ann in annotations if ann.get("modality") == "RGB"]
        ti_annotations = [ann for ann in annotations if ann.get("modality") == "TI"]
        
        detail = {
            "dir": root,
            "source": source_path,
            "rgb_count": len(rgb_annotations),
            "ti_count": len(ti_annotations),
            "fixed": False,
        }
        
        # 如果没有分离的必要，跳过
        if not rgb_annotations or not ti_annotations:
            detail["message"] = "无需分离，标注都属于同一模态"
            result["details"].append(detail)
            continue
            
        # 准备保存路径
        t_annotation_path = os.path.join(t_folder, "annotations.json") if t_folder else None
        v_annotation_path = os.path.join(v_folder, "annotations.json") if v_folder else None
        
        # 确保目标文件夹存在
        if ti_annotations and not t_folder:
            os.makedirs(os.path.join(root, "T"), exist_ok=True)
            t_folder = os.path.join(root, "T")
            t_annotation_path = os.path.join(t_folder, "annotations.json")
            
        if rgb_annotations and not v_folder:
            os.makedirs(os.path.join(root, "V"), exist_ok=True)
            v_folder = os.path.join(root, "V")
            v_annotation_path = os.path.join(v_folder, "annotations.json")
            
        # 保存RGB标注到V目录
        if rgb_annotations and v_annotation_path:
            rgb_data = {"annotations": rgb_annotations}
            try:
                with open(v_annotation_path, 'w', encoding='utf-8') as f:
                    json.dump(rgb_data, f, indent=2, ensure_ascii=False)
                detail["rgb_saved"] = True
            except Exception as e:
                detail["rgb_error"] = str(e)
                result["errors"].append(f"保存RGB标注失败: {v_annotation_path}, 错误: {str(e)}")
                
        # 保存TI标注到T目录
        if ti_annotations and t_annotation_path:
            ti_data = {"annotations": ti_annotations}
            try:
                with open(t_annotation_path, 'w', encoding='utf-8') as f:
                    json.dump(ti_data, f, indent=2, ensure_ascii=False)
                detail["ti_saved"] = True
            except Exception as e:
                detail["ti_error"] = str(e)
                result["errors"].append(f"保存TI标注失败: {t_annotation_path}, 错误: {str(e)}")
                
        # 如果source_annotations在根目录，可以考虑备份后删除
        if source_path != t_annotation_path and source_path != v_annotation_path:
            backup_path = source_path + ".bak"
            try:
                # 备份原始文件
                with open(backup_path, 'w', encoding='utf-8') as f:
                    json.dump(source_annotations, f, indent=2, ensure_ascii=False)
                detail["backed_up"] = True
                
                # 可选：删除原始文件
                # os.remove(source_path)
                # detail["removed_original"] = True
            except Exception as e:
                detail["backup_error"] = str(e)
                result["errors"].append(f"备份原始标注失败: {source_path}, 错误: {str(e)}")
        
        # 标记为已修复
        if detail.get("rgb_saved") or detail.get("ti_saved"):
            detail["fixed"] = True
            result["fixed"] += 1
            
        result["details"].append(detail)
    
    return result

def run_annotation_fixer_cli():
    """命令行入口函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='修复标注文件结构，将TI和RGB标注分离到各自目录')
    parser.add_argument('project_dir', help='项目根目录路径')
    parser.add_argument('--dry-run', action='store_true', help='仅检查不执行修改')
    
    args = parser.parse_args()
    
    print(f"开始处理项目: {args.project_dir}")
    result = fix_annotations_structure(args.project_dir)
    
    print(f"\n处理完成! 总计处理 {result['processed']} 个目录，修复 {result['fixed']} 个标注文件。")
    
    if result["errors"]:
        print(f"\n出现了 {len(result['errors'])} 个错误:")
        for err in result["errors"]:
            print(f"- {err}")
            
    print("\n详细信息:")
    for detail in result["details"]:
        print(f"- 目录: {detail['dir']}")
        print(f"  源文件: {detail['source']}")
        print(f"  RGB标注: {detail['rgb_count']}, TI标注: {detail['ti_count']}")
        if detail.get("fixed"):
            print(f"  状态: 已修复")
        else:
            print(f"  状态: {detail.get('message', '未修改')}")
        print("")

if __name__ == "__main__":
    run_annotation_fixer_cli()
